# -*- coding: utf-8 -*-

import os
import torch
# import torchvision
import cv2
import numpy as np
import SimpleITK as sitk
import matplotlib.pyplot as plt


class KaggleSegDataset(torch.utils.data.Dataset):
    def __init__(self, image_dir, label_dir):
        if not os.path.isdir(image_dir):
            raise RuntimeError("Image directory: "+ image_dir +" not found")
        
        if not os.path.isdir(label_dir):
            raise RuntimeError("Label directory: "+ label_dir +" not found")
        
        self.image_dir = image_dir
        self.label_dir = label_dir
        self.image_filelist = os.listdir(image_dir)
        self.label_filelist = os.listdir(label_dir)
        
        
        if len(self.image_filelist) != len(self.label_filelist):
            raise RuntimeError("Images amount in image directory and label directory should be identical.")
        
        self.length = len(self.image_filelist)      
    
    
    def load_dicom(self, filepath, label=False):
        image = sitk.GetArrayFromImage(sitk.ReadImage(filepath))[0]
        image = torch.from_numpy(image.astype(np.int32)).cuda()
        
        image = image.resize_((512, 512))
        image = torch.rot90(image, k=2)
        
        if not label:
            image = torch.unsqueeze(image, dim=0)
            # pass
        else:
            background = image == 0
            lung = (image == 1) + (image == 2)
            lesion = image == 3
            image = torch.stack((background, lung, lesion), axis=0)
            
            
        if label:
            image = torch.argmax(image.long(), dim=0)
        
        return image
    
    
    def __getitem__(self, index):
        if index >= self.length or index < 0:
            raise RuntimeError("index out of bounds.")
        
        image_path = os.path.join(self.image_dir, self.image_filelist[index])
        image = self.load_dicom(image_path)
        label_path = os.path.join(self.label_dir, self.label_filelist[index])
        label = self.load_dicom(label_path, True)

        image = (torch.clamp(image, min=-1200, max=600) + 1200) / 1800.0 * 255
        image = torch.unsqueeze((image.squeeze() * (label > 0) + 170 * (label == 0)), dim=0)
        
        
        return image, label
        
    
    def __len__(self):
        return self.length



class CCSegDataset(torch.utils.data.Dataset):
    def __init__(self, image_dir, label_dir):
        if not os.path.isdir(image_dir):
            raise RuntimeError("Image directory: "+ image_dir +" not found")
        
        if not os.path.isdir(label_dir):
            raise RuntimeError("Label directory: "+ label_dir +" not found")
            
        self.image_dir = image_dir
        self.label_dir = label_dir
        
        self.image_filelist = []
        self.label_filelist = []
        label_doclist = os.listdir(label_dir)
        for doc in label_doclist:
            image_docpath = os.path.join(image_dir, doc)
            label_docpath = os.path.join(label_dir, doc)
            if os.path.isdir(label_docpath):
                self.image_filelist += [os.path.join(image_docpath, x).replace("png", "jpg") for x in os.listdir(label_docpath)]
                self.label_filelist += [os.path.join(label_docpath, x) for x in os.listdir(label_docpath)]
                
        self.length = len(self.image_filelist)
        
    def __getitem__(self, index):
        if index >= self.length or index < 0:
            raise RuntimeError("index out of bounds.")
        
        image = cv2.imread(self.image_filelist[index])[..., 0]
        image = torch.from_numpy(image).cuda().float()
        image = torch.unsqueeze(image, dim=0) - 52
        image = torch.clamp(image, min=0, max=255) / (255 - 52) * 255
        
        label = cv2.imread(self.label_filelist[index])[..., 0]
        label = torch.from_numpy(label).cuda().long()
        
        image = torch.unsqueeze((image.squeeze() * (label > 0) + 170 * (label == 0)), dim=0)
        
        return image, label
        
    def __len__(self):
        return self.length



class MedSegDataset(torch.utils.data.Dataset):
    def __init__(self, image_path, mask_path):
        if not os.path.exists(image_path):
            raise RuntimeError("Image directory: "+ image_path +" not found")
        
        if not os.path.exists(mask_path):
            raise RuntimeError("Label directory: "+ mask_path +" not found")
            
        self.image_path = image_path
        self.mask_path = mask_path
        
        self.image_matrix = self.load_npy(image_path)
        self.mask_matrix = self.load_npy(mask_path, True)
    
    def load_npy(self, filepath, mask=False):
        image = np.load(filepath)
        # image = np.rot90(image, k=3, axes=(1, 2))
        image = image.transpose((0, 3, 2, 1))
        
        if mask:
            image = image[:, [3, 2, 0, 1], :, :]
        
        image = torch.from_numpy(image.copy()).cuda()
        if mask:
            image = torch.argmax(image.int(), dim=1)
        else:
            image = (torch.clamp(image, min=-1200, max=600) + 1200) / 1800.0
            image = image.float()
        
        return image
    
    def __getitem__(self, index):
        return self.image_matrix[index], self.mask_matrix[index]
    
    def __len__(self):
        return self.image_matrix.shape[0]


def visualize(image, label):
    plt.subplot(1, 2, 1)
    plt.imshow(image[0].cpu())
    plt.subplot(1, 2, 2)
    plt.imshow(label.cpu())
    plt.show()
    
# dataset = DcmSegDataset("D:\data\covid19-ct-scans\scans", "D:\data\covid19-ct-scans\masks")
# trainset, validset = torch.utils.data.random_split(dataset, [int(len(dataset)*0.9), len(dataset)-int(len(dataset)*0.9)])


# plt.imshow(validset[50][0][0])
# print(validset[1][1].shape)

# validLoader = torch.utils.data.DataLoader(validset, 3)

# for i, item in enumerate(validLoader):
#     if i >= 1:
#         break
#     x, y = item
#     print(y.shape)
#     plt.subplot(1, 2, 1)
#     plt.imshow(x[0])
#     plt.subplot(1, 2, 2)
#     plt.imshow(y[0])
#     plt.show()

# kaggle_dataset = KaggleSegDataset("D:\data\covid19-ct-scans\scans", "D:\data\covid19-ct-scans\masks")
# cc_dataset = CCSegDataset("D:\data\ct_lesion_seg\image", "D:\data\ct_lesion_seg\mask")
# # dataset = MedSegDataset("D:\data\covid-segmentation\images_radiopedia.npy", "D:\data\covid-segmentation\masks_radiopedia.npy")
# kaggle_loader = torch.utils.data.DataLoader(kaggle_dataset, batch_size=1, shuffle=True)
# cc_loader = torch.utils.data.DataLoader(cc_dataset, batch_size=1, shuffle=True)


# kaggle_tot = torch.zeros((1, 1, 512, 512)).cuda()
# for i, item in enumerate(kaggle_loader):
#     kaggle_tot += item[0]

# cc_tot = torch.zeros((1, 1, 512, 512)).cuda()
# for i, item in enumerate(cc_loader):
#     cc_tot += item[0]

# print(y.shape)
# visualize(x, y)

# x, y = dataset[400]
# visualize(x, y)
